#include <android/asset_manager_jni.h>
#include <android/bitmap.h>
#include <android/log.h>

#include <jni.h>

#include <string>
#include <vector>
#include <sstream>
// ncnn
#include "layer.h"
#include "net.h"
#include "benchmark.h"

static ncnn::UnlockedPoolAllocator g_blob_pool_allocator;
static ncnn::PoolAllocator g_workspace_pool_allocator;

static ncnn::Net yolov7;

class YoloV5Focus : public ncnn::Layer {
public:
    YoloV5Focus() {
        one_blob_only = true;
    }

    virtual int
    forward(const ncnn::Mat &bottom_blob, ncnn::Mat &top_blob, const ncnn::Option &opt) const {
        int w = bottom_blob.w;
        int h = bottom_blob.h;
        int channels = bottom_blob.c;

        int outw = w / 2;
        int outh = h / 2;
        int outc = channels * 4;

        top_blob.create(outw, outh, outc, 4u, 1, opt.blob_allocator);
        if (top_blob.empty())
            return -100;

#pragma omp parallel for num_threads(opt.num_threads)
        for (int p = 0; p < outc; p++) {
            const float *ptr = bottom_blob.channel(p % channels).row((p / channels) % 2) +
                               ((p / channels) / 2);
            float *outptr = top_blob.channel(p);

            for (int i = 0; i < outh; i++) {
                for (int j = 0; j < outw; j++) {
                    *outptr = *ptr;

                    outptr += 1;
                    ptr += 2;
                }

                ptr += w;
            }
        }

        return 0;
    }
};

DEFINE_LAYER_CREATOR(YoloV5Focus)

struct Object {
    float x;
    float y;
    float w;
    float h;
    int label;
    float prob;
};

std::string ObjectToString(const Object &obj) {
    std::ostringstream oss;
    oss << "Object("
        << "x=" << obj.x << ", "
        << "y=" << obj.y << ", "
        << "w=" << obj.w << ", "
        << "h=" << obj.h << ", "
        << "label=" << obj.label << ", "
        << "prob=" << obj.prob << ")";
    return oss.str();
}

std::string VectorToObjectString(const std::vector<Object> &objects) {
    std::ostringstream oss;
    for (size_t i = 0; i < objects.size(); ++i) {
        oss << ObjectToString(objects[i]);
        if (i < objects.size() - 1) {
            oss << " "; // 在对象之间添加空格分隔
        }
    }
    return oss.str();
}

static inline float intersection_area(const Object &a, const Object &b) {
    if (a.x > b.x + b.w || a.x + a.w < b.x || a.y > b.y + b.h || a.y + a.h < b.y) {
        // no intersection
        return 0.f;
    }

    float inter_width = std::min(a.x + a.w, b.x + b.w) - std::max(a.x, b.x);
    float inter_height = std::min(a.y + a.h, b.y + b.h) - std::max(a.y, b.y);

    return inter_width * inter_height;
}

static void qsort_descent_inplace(std::vector<Object> &faceobjects, int left, int right) {
    int i = left;
    int j = right;
    float p = faceobjects[(left + right) / 2].prob;

    while (i <= j) {
        while (faceobjects[i].prob > p)
            i++;

        while (faceobjects[j].prob < p)
            j--;

        if (i <= j) {
            // swap
            std::swap(faceobjects[i], faceobjects[j]);

            i++;
            j--;
        }
    }

#pragma omp parallel sections
    {
#pragma omp section
        {
            if (left < j) qsort_descent_inplace(faceobjects, left, j);
        }
#pragma omp section
        {
            if (i < right) qsort_descent_inplace(faceobjects, i, right);
        }
    }
}

static void qsort_descent_inplace(std::vector<Object> &faceobjects) {
    if (faceobjects.empty())
        return;

    qsort_descent_inplace(faceobjects, 0, faceobjects.size() - 1);
}

static void nms_sorted_bboxes(const std::vector<Object> &faceobjects, std::vector<int> &picked,
                              float nms_threshold) {
    picked.clear();

    const int n = faceobjects.size();

    std::vector<float> areas(n);
    for (int i = 0; i < n; i++) {
        areas[i] = faceobjects[i].w * faceobjects[i].h;
    }

    for (int i = 0; i < n; i++) {
        const Object &a = faceobjects[i];

        int keep = 1;
        for (int j = 0; j < (int) picked.size(); j++) {
            const Object &b = faceobjects[picked[j]];

            // intersection over union
            float inter_area = intersection_area(a, b);
            float union_area = areas[i] + areas[picked[j]] - inter_area;
            // float IoU = inter_area / union_area
            if (inter_area / union_area > nms_threshold)
                keep = 0;
        }

        if (keep)
            picked.push_back(i);
    }
}

static inline float sigmoid(float x) {
    return static_cast<float>(1.f / (1.f + exp(-x)));
}


static void generate_proposals(const ncnn::Mat &anchors, int stride, const ncnn::Mat &in_pad,
                               const ncnn::Mat &feat_blob, float prob_threshold,
                               std::vector<Object> &objects) {
    const int num_grid_x = feat_blob.w;
    const int num_grid_y = feat_blob.h;

    const int num_anchors = anchors.w / 2;

    const int num_class = feat_blob.c / num_anchors - 5;

    const int feat_offset = num_class + 5;

    for (int q = 0; q < num_anchors; q++) {
        const float anchor_w = anchors[q * 2];
        const float anchor_h = anchors[q * 2 + 1];

        for (int i = 0; i < num_grid_y; i++) {
            for (int j = 0; j < num_grid_x; j++) {
                // find class index with max class score
                int class_index = 0;
                float class_score = -FLT_MAX;
                for (int k = 0; k < num_class; k++) {
                    float score = feat_blob.channel(q * feat_offset + 5 + k).row(i)[j];
                    if (score > class_score) {
                        class_index = k;
                        class_score = score;
                    }
                }

                float box_score = feat_blob.channel(q * feat_offset + 4).row(i)[j];

                float confidence = sigmoid(box_score) * sigmoid(class_score);

                if (confidence >= prob_threshold) {
                    float dx = sigmoid(feat_blob.channel(q * feat_offset + 0).row(i)[j]);
                    float dy = sigmoid(feat_blob.channel(q * feat_offset + 1).row(i)[j]);
                    float dw = sigmoid(feat_blob.channel(q * feat_offset + 2).row(i)[j]);
                    float dh = sigmoid(feat_blob.channel(q * feat_offset + 3).row(i)[j]);

                    float pb_cx = (dx * 2.f - 0.5f + j) * stride;
                    float pb_cy = (dy * 2.f - 0.5f + i) * stride;

                    float pb_w = pow(dw * 2.f, 2) * anchor_w;
                    float pb_h = pow(dh * 2.f, 2) * anchor_h;

                    float x0 = pb_cx - pb_w * 0.5f;
                    float y0 = pb_cy - pb_h * 0.5f;
                    float x1 = pb_cx + pb_w * 0.5f;
                    float y1 = pb_cy + pb_h * 0.5f;

                    Object obj;
                    obj.x = x0;
                    obj.y = y0;
                    obj.w = x1 - x0;
                    obj.h = y1 - y0;
                    obj.label = class_index;
                    obj.prob = confidence;

                    objects.push_back(obj);
                }
            }
        }
    }
}


extern "C" {

// FIXME DeleteGlobalRef is missing for objCls
static jclass objCls = NULL;
static jmethodID constructortorId;
static jfieldID xId;
static jfieldID yId;
static jfieldID wId;
static jfieldID hId;
static jfieldID labelId;
static jfieldID probId;

JNIEXPORT jint JNI_OnLoad(JavaVM *vm, void *reserved) {
    __android_log_print(ANDROID_LOG_DEBUG, "YoloV7Ncnn", "JNI_OnLoad");
    ncnn::create_gpu_instance();
    return JNI_VERSION_1_4;
}

JNIEXPORT void JNI_OnUnload(JavaVM *vm, void *reserved) {
    __android_log_print(ANDROID_LOG_DEBUG, "YoloV7Ncnn", "JNI_OnUnload");

    ncnn::destroy_gpu_instance();
}

JNIEXPORT jboolean JNICALL
Java_com_tencent_ncnnyolov7_NcnnYolov7_InitUSE(JNIEnv *env, jobject thiz, jobject assetManager) {
    yolov7.clear();
    ncnn::Option opt;
    opt.lightmode = true;
    opt.num_threads = 4;
    opt.blob_allocator = &g_blob_pool_allocator;
    opt.workspace_allocator = &g_workspace_pool_allocator;
    opt.use_packing_layout = true;

    if (ncnn::get_gpu_count() != 0)
        opt.use_vulkan_compute = true;

    AAssetManager *mgr = AAssetManager_fromJava(env, assetManager);

    yolov7.opt = opt;

    // init param
    {
        int ret = yolov7.load_param(mgr, "yolov7.param");
        if (ret != 0) {
            __android_log_print(ANDROID_LOG_DEBUG, "YoloV7Ncnn", "load_param failed");
            return JNI_FALSE;
        }
    }

    // init bin
    {
        int ret = yolov7.load_model(mgr, "yolov7.bin");
        if (ret != 0) {
            __android_log_print(ANDROID_LOG_DEBUG, "YoloV7Ncnn", "load_model failed");
            return JNI_FALSE;
        }
    }

    // init jni glue
    jclass localObjCls = env->FindClass("com/tencent/ncnnyolov7/NcnnYolov7$Obj");
    objCls = reinterpret_cast<jclass>(env->NewGlobalRef(localObjCls));

    constructortorId = env->GetMethodID(objCls, "<init>", "(Lcom/tencent/ncnnyolov7/NcnnYolov7;)V");

    xId = env->GetFieldID(objCls, "x", "F");
    yId = env->GetFieldID(objCls, "y", "F");
    wId = env->GetFieldID(objCls, "w", "F");
    hId = env->GetFieldID(objCls, "h", "F");
    labelId = env->GetFieldID(objCls, "label", "Ljava/lang/String;");
    probId = env->GetFieldID(objCls, "prob", "F");

    return JNI_TRUE;
}


JNIEXPORT jboolean JNICALL
Java_com_tencent_ncnnyolov7_NcnnYolov7_InitCrystal(JNIEnv *env, jobject thiz,
                                                   jobject assetManager) {
    yolov7.clear();
    ncnn::Option opt;
    opt.lightmode = true;
    opt.num_threads = 4;
    opt.blob_allocator = &g_blob_pool_allocator;
    opt.workspace_allocator = &g_workspace_pool_allocator;
    opt.use_packing_layout = true;

    if (ncnn::get_gpu_count() != 0)
        opt.use_vulkan_compute = true;

    AAssetManager *mgr = AAssetManager_fromJava(env, assetManager);

    yolov7.opt = opt;

    {
        int ret = yolov7.load_param(mgr, "crystal.ncnn.param");
        if (ret != 0) {
            __android_log_print(ANDROID_LOG_DEBUG, "YoloV7Ncnn", "load_param failed");
            return JNI_FALSE;
        }
    }

    {
        int ret = yolov7.load_model(mgr, "crystal.ncnn.bin");
        if (ret != 0) {
            __android_log_print(ANDROID_LOG_DEBUG, "YoloV7Ncnn", "load_model failed");
            return JNI_FALSE;
        }
    }

    jclass localObjCls = env->FindClass("com/tencent/ncnnyolov7/NcnnYolov7$Obj");
    objCls = reinterpret_cast<jclass>(env->NewGlobalRef(localObjCls));

    constructortorId = env->GetMethodID(objCls, "<init>", "(Lcom/tencent/ncnnyolov7/NcnnYolov7;)V");

    xId = env->GetFieldID(objCls, "x", "F");
    yId = env->GetFieldID(objCls, "y", "F");
    wId = env->GetFieldID(objCls, "w", "F");
    hId = env->GetFieldID(objCls, "h", "F");
    labelId = env->GetFieldID(objCls, "label", "Ljava/lang/String;");
    probId = env->GetFieldID(objCls, "prob", "F");

    return JNI_TRUE;
}

// public native Obj[] Detect(Bitmap bitmap, boolean use_gpu);
JNIEXPORT jobjectArray JNICALL
Java_com_tencent_ncnnyolov7_NcnnYolov7_Detect(JNIEnv *env, jobject thiz, jobject bitmap,
                                              jboolean use_gpu, jint model) {
    __android_log_print(ANDROID_LOG_DEBUG, "YoloV7Ncnn", "start detect");

    if (use_gpu == JNI_TRUE && ncnn::get_gpu_count() == 0) {
        return NULL;
    }

    double start_time = ncnn::get_current_time();

    AndroidBitmapInfo info;
    AndroidBitmap_getInfo(env, bitmap, &info);
    const int width = info.width;
    const int height = info.height;
    if (info.format != ANDROID_BITMAP_FORMAT_RGBA_8888)
        return NULL;

    // ncnn from bitmap
    const int target_size = 640;

    // letterbox pad to multiple of 32
    int w = width;
    int h = height;
    float scale = 1.f;
    if (w > h) {
        scale = (float) target_size / w;
        w = target_size;
        h = h * scale;
    } else {
        scale = (float) target_size / h;
        h = target_size;
        w = w * scale;
    }

    ncnn::Mat in = ncnn::Mat::from_android_bitmap_resize(env, bitmap, ncnn::Mat::PIXEL_RGB, w, h);

    int wpad = (w + 31) / 32 * 32 - w;
    int hpad = (h + 31) / 32 * 32 - h;
    ncnn::Mat in_pad;
    ncnn::copy_make_border(in, in_pad, hpad / 2, hpad - hpad / 2, wpad / 2, wpad - wpad / 2,
                           ncnn::BORDER_CONSTANT, 114.f);

    std::vector<Object> objects;
    {
        const float prob_threshold = 0.25f;
        const float nms_threshold = 0.45f;
        const float norm_vals[3] = {1 / 255.f, 1 / 255.f, 1 / 255.f};
        in_pad.substract_mean_normalize(0, norm_vals);
        ncnn::Extractor ex = yolov7.create_extractor();
        ex.set_vulkan_compute(use_gpu);
        ex.input("in0", in_pad);
        std::vector<Object> proposals;

        // stride 8
        {
            ncnn::Mat out;
            ex.extract("out0", out);

            ncnn::Mat anchors(6);
            anchors[0] = 12.f;
            anchors[1] = 16.f;
            anchors[2] = 19.f;
            anchors[3] = 36.f;
            anchors[4] = 40.f;
            anchors[5] = 28.f;

            std::vector<Object> objects8;
            generate_proposals(anchors, 8, in_pad, out, prob_threshold, objects8);
            proposals.insert(proposals.end(), objects8.begin(), objects8.end());
        }

        // stride 16
        {
            ncnn::Mat out;
            ex.extract("out1", out);

            ncnn::Mat anchors(6);
            anchors[0] = 36.f;
            anchors[1] = 75.f;
            anchors[2] = 76.f;
            anchors[3] = 55.f;
            anchors[4] = 72.f;
            anchors[5] = 146.f;

            std::vector<Object> objects16;
            generate_proposals(anchors, 16, in_pad, out, prob_threshold, objects16);
            proposals.insert(proposals.end(), objects16.begin(), objects16.end());
        }

        // stride 32
        {
            ncnn::Mat out;
            ex.extract("out2", out);

            ncnn::Mat anchors(6);
            anchors[0] = 142.f;
            anchors[1] = 110.f;
            anchors[2] = 192.f;
            anchors[3] = 243.f;
            anchors[4] = 459.f;
            anchors[5] = 401.f;

            std::vector<Object> objects32;
            generate_proposals(anchors, 32, in_pad, out, prob_threshold, objects32);

            proposals.insert(proposals.end(), objects32.begin(), objects32.end());
        }

        // sort all proposals by score from highest to lowest
        qsort_descent_inplace(proposals);
        // apply nms with nms_threshold
        std::vector<int> picked;
        nms_sorted_bboxes(proposals, picked, nms_threshold);

        int count = picked.size();

        objects.resize(count);
        for (int i = 0; i < count; i++) {
            objects[i] = proposals[picked[i]];
            float x0 = (objects[i].x - (wpad / 2)) / scale;
            float y0 = (objects[i].y - (hpad / 2)) / scale;
            float x1 = (objects[i].x + objects[i].w - (wpad / 2)) / scale;
            float y1 = (objects[i].y + objects[i].h - (hpad / 2)) / scale;

            // clip
            x0 = std::max(std::min(x0, (float) (width - 1)), 0.f);
            y0 = std::max(std::min(y0, (float) (height - 1)), 0.f);
            x1 = std::max(std::min(x1, (float) (width - 1)), 0.f);
            y1 = std::max(std::min(y1, (float) (height - 1)), 0.f);

            objects[i].x = x0;
            objects[i].y = y0;
            objects[i].w = x1 - x0;
            objects[i].h = y1 - y0;
        }
    }

    if (model == 0) {
        static const char *class_names[] = {
                "管型",
                "晶体",
                "上皮细胞",
                "上皮细胞核",
                "红细胞",
                "白细胞",
                "霉菌"
        };
        jobjectArray jObjArray = env->NewObjectArray(objects.size(), objCls, NULL);

        for (size_t i = 0; i < objects.size(); i++) {
            jobject jObj = env->NewObject(objCls, constructortorId, thiz);
            env->SetFloatField(jObj, xId, objects[i].x);
            env->SetFloatField(jObj, yId, objects[i].y);
            env->SetFloatField(jObj, wId, objects[i].w);
            env->SetFloatField(jObj, hId, objects[i].h);
            __android_log_print(ANDROID_LOG_DEBUG, "YoloV7Ncnn", "检测结果为：%i", objects[i].label);
            env->SetObjectField(jObj, labelId, env->NewStringUTF(class_names[objects[i].label]));
            env->SetFloatField(jObj, probId, objects[i].prob);

            env->SetObjectArrayElement(jObjArray, i, jObj);
        }

        double elasped = ncnn::get_current_time() - start_time;
        __android_log_print(ANDROID_LOG_DEBUG, "YoloV7Ncnn", "%.2fms   detect", elasped);
        __android_log_print(ANDROID_LOG_DEBUG, "YoloV7Ncnn", "类别：%s",
                            VectorToObjectString(objects).c_str());
        return jObjArray;
    } else if (model == 1) {
        static const char *class_names[] = {
                "一水草酸钙",
                "二水草酸钙",
                "尿酸",
                "磷酸铵镁",
                "胱氨酸"
        };
        jobjectArray jObjArray = env->NewObjectArray(objects.size(), objCls, NULL);

        for (size_t i = 0; i < objects.size(); i++) {
            jobject jObj = env->NewObject(objCls, constructortorId, thiz);
            env->SetFloatField(jObj, xId, objects[i].x);
            env->SetFloatField(jObj, yId, objects[i].y);
            env->SetFloatField(jObj, wId, objects[i].w);
            env->SetFloatField(jObj, hId, objects[i].h);
            __android_log_print(ANDROID_LOG_DEBUG, "YoloV7Ncnn", "检测结果为：%i", objects[i].label);
            env->SetObjectField(jObj, labelId, env->NewStringUTF(class_names[objects[i].label]));
            env->SetFloatField(jObj, probId, objects[i].prob);

            env->SetObjectArrayElement(jObjArray, i, jObj);
        }

        double elasped = ncnn::get_current_time() - start_time;
        __android_log_print(ANDROID_LOG_DEBUG, "YoloV7Ncnn", "%.2fms   detect", elasped);
        __android_log_print(ANDROID_LOG_DEBUG, "YoloV7Ncnn", "类别：%s",
                            VectorToObjectString(objects).c_str());
        return jObjArray;
    }
}

}